// Copyright (c) 2023 Presto Labs Pte. Ltd.
// Author: soony

#include <cxxopts.hpp>

#include "coin/proto/coin_market_enums.pb.h"
#include "coin2/app/driver.h"
#include "coin2/base/config/app_config.h"
#include "coin2/base/config/args.h"
#include "coin2/base/feed_convert_fastfeature.h"
#include "coin2/base/log.h"
#include "coin2/base/time.h"
#include "coin2/exchange/base/feed/update.h"
#include "coin2/strategy/linear_model/product_strategy.h"
#include "coin2/strategy/strategy.h"
#include "presto/quant/base/init.h"
#include "presto/quant/base/module.h"
#include "presto/quant/feature/feature.h"
#include "presto/quant/feature/feature_manager.h"
#include "presto/quant/feature/feature_dumper2.h"
#include "presto/quant/symbology/symbol.h"
#include "presto/quant/symbology/symbol_manager.h"

using coin2::app::Driver;
using coin2::base::ConvertFeedMessage;
using coin2::exchange::base::symbology::ProductEncyclopedia;
using coin2::exchange::base::feed::FeedUpdate;
// using coin2::feed::fastfeed::FastFeedUpdate;
using coin2::exchange::base::market::MarketExchangeApi;
using coin2::exchange::base::market::ExchangeTypeToString;
using coin2::exchange::base::market::MarketTypeToString;
using coin2::strategy::linear_model::FeatureFeedConverter;
using coin2::strategy::linear_model::GetSymbolInline;

using ::presto::Symbol;
using ::presto::SymbolManager;
using ::presto::quant::feature::Dumper2;
using ::presto::quant::feature::FeatureManager;
using ::util::json::ListValue;
using ::util::json::DictionaryValue;

inline double flip_pricer(double price1, double price2) {
  return price1 / price2 * price1;
}

class ReverseBookBuilder : public BookBuilder {

 public:
  explicit ReverseBookBuilder(bool enable_diff_book) : BookBuilder(enable_diff_book) {}

  ReverseBookBuilder(const BookBuilder& from_bb, size_t max_book_level)
    : BookBuilder(from_bb, max_book_level) {}

  void BuildReverseBook(const BookBuilder* from_bb, double price) {
    asks_.clear();
    bids_.clear();
    double flip_price = 0.0;
    double flip_qty = 0.0;
    for (auto&& [prc, qty] : from_bb->Bids()) {
      // TODO: adjust tick price -> merge qty
      flip_price = flip_pricer(price, prc);
      flip_qty = prc / flip_price * qty;
      asks_[flip_price] = flip_qty;
    }
    for (auto&& [prc, qty] : from_bb->Asks()) {
      flip_price = flip_pricer(price, prc);
      flip_qty = prc / flip_price * qty;
      bids_[flip_price] = flip_qty;
    }
  }
  // ~ReverseBookBuilder() { CleanUp(tin_ask_price_, max_bid_price_); }
  ~ReverseBookBuilder() { }
};

class FeatureFeedAugmentor {
 public:
  virtual ~FeatureFeedAugmentor() = default;

  virtual const FeedUpdate& onTradeFeed(const FeedUpdate& upd) = 0;
  virtual const FeedUpdate& onBookFeed(const FeedUpdate& upd) = 0;
  // virtual FeedUpdate& onStatusFeed(const FeedUpdate& upd);

 public:
  const FeatureManager* feature_manager;
  const SymbolManager* symbol_manager;

 private:
  const ProductEncyclopedia* product_cache_;
};


class FlipAugmentor : public FeatureFeedAugmentor {
 public:
  explicit FlipAugmentor(
      const FeatureManager* feature_manager,
      const ProductEncyclopedia* product_cache,
      FeatureFeedConverter* converter)
      : feature_manager(feature_manager),
        symbol_manager(feature_manager->GetSymbolManager()),
        product_cache_(product_cache),
        converter_(converter) {

    LOG(INFO) << "Initialize FlipAugmentor";
    const auto& fspec_pb_ = feature_manager->GetFspecPb();
    for (const auto& feature_instance : fspec_pb_->features()) {
      std::string jsonstr;
      CHECK(::google::protobuf::util::MessageToJsonString(feature_instance.config_obj(), &jsonstr).ok())
          << feature_instance.config_obj().DebugString();
      nlohmann::json jf = nlohmann::json::parse(jsonstr);
      CHECK(jf.is_object());
      DictionaryValue* config_obj(reinterpret_cast<DictionaryValue*>(new nlohmann::json(jf)));
      if (!config_obj->HasKey("cross_list")) {
        continue;
      }

      ListValue* cross_list;
      CHECK(config_obj->GetList("cross_list", &cross_list)) << feature_instance.config_obj().DebugString();
      for (int i = 0; i < cross_list->GetSize(); ++i) {
        DictionaryValue* cross_dict;
        CHECK(cross_list->GetDictionary(i, &cross_dict));
        std::string exchange_str;
        std::string product_name;
        CHECK(cross_dict->GetString("focus_exchange", &exchange_str));
        CHECK(cross_dict->GetString("focus_product_name", &product_name));
        const Symbol* fcs = symbol_manager->FromString(exchange_str, product_name);
        CHECK(cross_dict->GetString("exchange", &exchange_str));
        CHECK(cross_dict->GetString("product_name", &product_name));
        const Symbol* ref = symbol_manager->FromString(exchange_str, product_name);
        double multiplier;
        CHECK(cross_dict->GetReal("multiplier", &multiplier));

        focus_symbol_[fcs] = fcs;
        focus_symbol_[ref] = fcs;
        multiplier_[fcs] = 1.0;
        multiplier_[ref] = multiplier;
      }
    }
  }

  FeedUpdate& CopyFeedUpdate(const FeedUpdate& upd) {
    // TODO: need to consider system_name ?
    if (update_ptr_.find(upd.mea()) == update_ptr_.end()) {
      update_ptr_[upd.mea()].reset(new FeedUpdate(upd.system_name(), upd.mea()));
    }

    FeedUpdate& update_ = *update_ptr_[upd.mea()].get();
    update_.Reset();
    update_.SetTimestamp(upd.timestamp());
    update_.SetRecipe(upd.recipe());
    update_.SetFeedType(upd.feed_type());
    update_.SetNativeSymbol(upd.native_symbol());
    update_.SetProduct(upd.product());

    return update_;
  }

  const FeedUpdate& onTradeFeed(const FeedUpdate& upd) {
    const Symbol* symbol = converter_->GetSymbol(upd);
    if (symbol == nullptr) {
      return upd;
    }
    const auto focus_symbol = focus_symbol_[symbol];
    auto trade = upd.trade();
    if (!first_pricer_.count(focus_symbol)) {
      first_pricer_[focus_symbol] = trade.price * multiplier_[symbol];
    }
    coin2::exchange::feed::TradeSide new_side;
    if (trade.side == coin2::exchange::feed::TradeSide::TRADE_BUY_SIDE) {
      new_side = coin2::exchange::feed::TradeSide::TRADE_SELL_SIDE;
    } else if (trade.side == coin2::exchange::feed::TradeSide::TRADE_SELL_SIDE) {
      new_side = coin2::exchange::feed::TradeSide::TRADE_BUY_SIDE;
    } else {
      throw std::runtime_error("invalid trade side");
    }
    const auto flip_price = flip_pricer(first_pricer_.at(focus_symbol) / multiplier_[symbol], trade.price);
    const auto flip_qty = trade.price / flip_price * trade.fill_qty;
    trade_ = {
        trade.timestamp,
        flip_price,
        flip_qty,
        new_side,
        trade.has_more_trade
    };
    auto& new_upd = CopyFeedUpdate(upd);
    new_upd.SetTrade(&trade_);
    new_upd.MutableTs().set_exchange_publish_timestamp(trade.timestamp);
    return new_upd;
  }

  const FeedUpdate& onBookFeed(const FeedUpdate& upd) {
    const Symbol* symbol = converter_->GetSymbol(upd);
    if (symbol == nullptr) {
      return upd;
    }
    const auto focus_symbol = focus_symbol_[symbol];

    if (!first_pricer_.count(focus_symbol)) {
      first_pricer_[focus_symbol] = upd.GetBookBuilder()->Bid0()->price * multiplier_[symbol];
    }

    ReverseBookBuilder* bb = books_[symbol].get();
    if (bb == nullptr) {
      books_[symbol].reset(bb = new ReverseBookBuilder(false));
    }
    bb->BuildReverseBook(upd.GetBookBuilder(), first_pricer_.at(focus_symbol) / multiplier_[symbol]);
    auto& new_upd = CopyFeedUpdate(upd);
    new_upd.SetBook(bb);
    new_upd.MutableTs().set_exchange_publish_timestamp(bb->Timestamp());
    return new_upd;
  }


 public:
  const FeatureManager* feature_manager;
  const SymbolManager* symbol_manager;

 private:
  presto::detail::FastSymbolMap<std::unique_ptr<ReverseBookBuilder>> books_;
  Trade trade_{};

  std::unordered_map<MarketExchangeApi, std::unique_ptr<FeedUpdate>> update_ptr_;
  presto::detail::FastSymbolMap<double> first_pricer_;
  presto::detail::FastSymbolMap<const Symbol*> focus_symbol_;
  presto::detail::FastSymbolMap<double> multiplier_;
  const ProductEncyclopedia* product_cache_;
  FeatureFeedConverter* converter_;
};


class FastFeatureAugmentationDumper : public coin2::strategy::IStrategy {
 public:
  FastFeatureAugmentationDumper(
      const std::string& sampler_name_subset,
      const std::experimental::filesystem::path& feature_info_file,
      const std::experimental::filesystem::path& recorder_config_file,
      const std::experimental::filesystem::path& kline_json_file)
      : dumper(new Dumper2(sampler_name_subset, feature_info_file,
            recorder_config_file, kline_json_file)) {
    LOG(INFO) << feature_info_file;
    LOG(INFO) << recorder_config_file;
    converter.reset(new FeatureFeedConverter(
        dumper->GetFeatureManager()->GetSymbolManager(),
        &product_cache,
        true));
    augmentor.reset(new FlipAugmentor(
        dumper->GetFeatureManager(),
        &product_cache,
        converter.get()));
  }

  ~FastFeatureAugmentationDumper() {
    LOG(INFO) << "Bye";
    dumper.reset();
  }

  void Init(::coin2::app::Driver* driver) override {}

  void onTradeFeed(const FeedUpdate& upd) override {
    const auto& new_upd = augmentor->onTradeFeed(upd);
    converter->onTradeFeed(new_upd);
    if (converter->feed_converted) {
      onFeatureFeed(new_upd, &converter->last_feed_msg);
    }
  }

  void onBookFeed(const FeedUpdate& upd) override {
    const auto& new_upd = augmentor->onBookFeed(upd);
    converter->onBookFeed(new_upd);
    if (converter->feed_converted) {
      onFeatureFeed(new_upd, &converter->last_feed_msg);
    }
  }

  void onStatusFeed(const FeedUpdate& upd) {
    converter->onStatusFeed(upd);
    if (converter->feed_converted) {
      // TODO: add onStatusFeed
      // const auto& new_upd = augmentor->onStatusFeed(upd);
      onFeatureFeed(upd, &converter->last_feed_msg);
    }
  }

  void onIndexFeed(const FeedUpdate& upd) override {
    onStatusFeed(upd);
  }

  void onLiquidationFeed(const FeedUpdate& upd) override {
    onStatusFeed(upd);
  }

  void onFeatureFeed(const FeedUpdate& upd, const ::fastfeature::FeedMessage* feed_msg) {
    const Symbol* symbol = converter->GetSymbol(upd);
    // this feature dumper won't take this symbol.
    if (symbol == nullptr) return;
    dumper->UpdateFeed(symbol, *feed_msg);
  }

 private:
  std::unique_ptr<Dumper2> dumper;
  std::unique_ptr<FeatureFeedConverter> converter;
  std::unique_ptr<FeatureFeedAugmentor> augmentor;
  ProductEncyclopedia product_cache;
};

int main(int argc, char* argv[]) {
  InitModules();

  coin2::app::DriverConfig app;

  cxxopts::Options opt("feature dumper", "feature dumper");
  opt.add_options()("sampler_name_subset", "", cxxopts::value<std::string>());
  opt.add_options()("config_filename", "feature config path", cxxopts::value<std::string>());
  opt.add_options()("recorder_config_filename", "", cxxopts::value<std::string>());
  opt.add_options()("kline_json", "kline info path",
      cxxopts::value<std::string>()->default_value(""));
  opt.add_options()("hours_warmedup", "hours_warmup that is already reflected",
      cxxopts::value<double>()->default_value("0.0"));
  coin2::base::config::AddDriverOptions(&opt);
  coin2::base::config::AddIntervalOptions(&opt);
  coin2::base::config::AddProductOptions(&opt);
  coin2::base::config::AddFeedArchiveOptions(&opt);

  auto res = opt.parse(argc, argv);
  coin2::base::config::ParseDriverOptions(res, &app);
  coin2::base::config::ParseIntervalOptions(
      res,
      app.mutable_feed()->mutable_common()->mutable_archive()->mutable_interval());
  coin2::base::config::ParseProductOptions(res, app.mutable_feed());
  app.mutable_feed()->mutable_common()->mutable_archive()->set_machine(
      res["feed_machine"].as<std::string>());
  LOG(INFO) << app.DebugString();

  ::coin2::base::InitLogging(argv[0], app.log());

  FastFeatureAugmentationDumper fastfeature_dumper(
      res["sampler_name_subset"].as<std::string>(),
      res["config_filename"].as<std::string>(),
      res["recorder_config_filename"].as<std::string>(),
      res["kline_json"].as<std::string>());
  {
    Driver driver(app, &fastfeature_dumper);
    driver.Run();
  }

  return 0;
}
